from langchain_community.llms import Ollama
from langchain_core.messages import HumanMessage, AIMessage
from langchain_core.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, \
    HumanMessagePromptTemplate, MessagesPlaceholder

llm = Ollama(model="llama3", base_url="http://localhost:11434")

def chat_with_ollama():
    chat_template = ChatPromptTemplate.from_messages([
        MessagesPlaceholder(variable_name="conversation"),
        HumanMessagePromptTemplate.from_template("请将上面的对话翻译为{language}。")
    ])

    human_message = HumanMessage(content="Who is Elon Musk?")
    ai_message = AIMessage(content="Elon Musk is a billionaire entrepreneur, inventor, and industrial designer")

    prompt = chat_template.format_prompt(language="中文",conversation=[human_message,ai_message])

    print(prompt.to_messages())
    response = llm.invoke(prompt)
    return response
if __name__ == '__main__':
    print(chat_with_ollama())